import os
import yaml
import torch
import typing
from PIL import Image
import interface.utils as Iutils
from torchvision import transforms, models


def predict_class(
        model: torch.nn.Module, img_path: str,
        class_names: typing.Dict[int, str], device: torch.device,
        image_size: int) -> str:
    """
    Predict the class of an image using the model.

    Args:
            model (torch.nn.Module): Trained model.
            img_path (str): Path to the input image.
            class_names (Dict[int, str]): Dictionary of class names.
            device (torch.device): Device for computation.
            image_size (int): Size of the input image.

    Returns:
            str: Predicted class name.
    """
    # Set the model to evaluation mode
    model.eval()

    # Open the image and apply transformations
    data_transforms = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
                             0.229, 0.224, 0.225])
    ])

    img = Image.open(img_path).convert("RGB")
    img = data_transforms(img)
    img = img.unsqueeze(0)
    img = img.to(device)

    # Prediction
    with torch.no_grad():
        outputs = model(img)
        _, preds = torch.max(outputs, 1)

    # Return the predicted class name
    return class_names[preds[0].item()]


def start():
    config = Iutils.load_config('test_config.yaml')
    class_config_path = config['class_config_path']

    class_names = Iutils.load_config(class_config_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = models.convnext_base(pretrained=False)
    num_classes = len(class_names)
    model.classifier[2] = torch.nn.Linear(
        model.classifier[2].in_features, num_classes)
    if device.type == 'cpu':
        model.load_state_dict(torch.load(
            config['model_weights_path'], map_location=torch.device('cpu')))
    elif device.type == 'cuda':
        model.load_state_dict(torch.load(config['model_weights_path']))
    model = model.to(device)

    predicted_class = predict_class(
        model, config['image_path'], class_names, device, config['image_size'])

    print(f'Predicted class: {predicted_class}')
    return predicted_class
